/**
 * Gooogle Gemini Provider
 * 文档：https://ai.google.dev/gemini-api/docs?hl=zh-cn
 */
import type { LLMCompletionsData, LLMCompletionsCb, LLMAttachCheckRes } from '@/provider/llm/type'
import { llmBaseProvider } from '@/provider/llm/baseProvider'
import { GoogleGenAI, type Content, type Part } from '@google/genai'
import { geneTitlePrompt, clearContent, extractBase64Data } from '@/provider/llm/hooks/utils'

class GeminiProvider extends llmBaseProvider {
  // 客户端
  private static client: GoogleGenAI
  // stream signal
  private static signalController: AbortController
  // 是否已打断
  private static isCanceled = false

  protected init() {
    // 初始化 openai-node
    GeminiProvider.client = new GoogleGenAI({
      apiKey: this.getApiKey(),
    })
  }

  public async accessCheck() {
    // 尝试调用 api 以判断是否可用
    try {
      await GeminiProvider.client.models.list({
        config: {
          pageSize: 1,
        },
      })
      return true
    } catch {
      return false
    }
  }

  public attachCheck(model: LLMModel): LLMAttachCheckRes {
    return {
      allowVision: !!model.vision,
      visionLimit: 20 * 1024 * 1024,
      allowDoc: true,
      docLimit: 20 * 1024 * 1024,
      docsAccept: ['application/pdf', 'text/plain'],
    }
  }

  public async getModels() {
    // Gemini 模型可以通过 api 获取全部，但数量过多，因此建议通过 defaultModels 进行配置选用
    // 可用模型参考：https://ai.google.dev/gemini-api/docs/models?hl=zh-cn
    return this.getDefaultModels()
  }

  public async geneTitle(model: string, messages: LLMMessage[]) {
    try {
      const result = await GeminiProvider.client.models.generateContent({
        model,
        contents: geneTitlePrompt(messages),
        config: {
          temperature: 0.3,
          topP: 0.5,
          // maxOutputTokens: 100
        },
      })

      return result.text
        .replace(/<[^>]*>/g, '')
        .replace(/\n/g, '')
        .trim()
    } catch (error) {
      window.console.error(error)
      return ''
    }
  }

  public async completions(data: LLMCompletionsData, callback: LLMCompletionsCb) {
    let fullText = ''
    const history: Content[] = []
    for (const item of data.messages) {
      if (item.role === 'user') {
        const parts: Part[] = []
        // 处理附件请求
        // 此处直接使用文件base64进行提交，仅作演示使用，生产环境下建议使用存储服务转储
        // 参考：https://ai.google.dev/gemini-api/docs/image-understanding?hl=zh-cn#upload-image
        if (item.attachs?.length) {
          for (const attach of item.attachs) {
            parts.push({
              inlineData: {
                mimeType: attach.mime,
                data: extractBase64Data(attach.url),
              },
            })
          }
        }
        // 添加文本
        parts.push({ text: clearContent(item.content) })
        history.push({
          role: 'user',
          parts,
        })
      } else if (item.role !== 'error') {
        history.push({
          role: 'model',
          parts: [{ text: clearContent(item.content) }],
        })
      }
    }
    // 重置打断信号
    GeminiProvider.signalController = new AbortController()
    GeminiProvider.isCanceled = false
    try {
      const chat = GeminiProvider.client.chats.create({
        model: data.model,
        // 忽略最后一条 history 数据
        history: history.slice(0, -1),
        config: {
          temperature: data.temperature,
          maxOutputTokens: data.max_tokens || undefined,
          abortSignal: GeminiProvider.signalController.signal,
        },
      })
      const stream = await chat.sendMessageStream({ message: history[history.length - 1].parts })
      for await (const chunk of stream) {
        if (GeminiProvider.isCanceled) {
          window.console.warn('Abort stream!')
          break
        }
        fullText += chunk.text
        callback.onUpdate(fullText)
      }
      callback.onSuccess(fullText)
    } catch (error) {
      window.console.error(error)
      if (error.name === 'AbortError') {
        // 打断
        callback.onSuccess(fullText)
      } else {
        callback.onError(error)
      }
    } finally {
      GeminiProvider.isCanceled = false
    }
  }

  public async abort() {
    GeminiProvider.signalController?.abort()
    GeminiProvider.isCanceled = true
  }
}

export default GeminiProvider
